package com.zzw.oauth2.conf;

import com.zzw.entity.OauthAccessToken;
import com.zzw.entity.OauthRefreshToken;
import com.zzw.service.OauthAccessTokenService;
import com.zzw.service.OauthRefreshTokenService;
import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.dao.EmptyResultDataAccessException;
import org.springframework.security.oauth2.common.OAuth2AccessToken;
import org.springframework.security.oauth2.common.OAuth2RefreshToken;
import org.springframework.security.oauth2.common.util.SerializationUtils;
import org.springframework.security.oauth2.provider.OAuth2Authentication;
import org.springframework.security.oauth2.provider.token.AuthenticationKeyGenerator;
import org.springframework.security.oauth2.provider.token.DefaultAuthenticationKeyGenerator;
import org.springframework.security.oauth2.provider.token.TokenStore;
import org.springframework.security.oauth2.provider.token.store.JdbcTokenStore;

import java.io.UnsupportedEncodingException;
import java.math.BigInteger;
import java.security.MessageDigest;
import java.security.NoSuchAlgorithmException;

import java.util.*;

public class MybatisTokenStore implements TokenStore {
    private static final Log LOG = LogFactory.getLog(JdbcTokenStore.class);

    @Autowired
    private OauthAccessTokenService oauthAccessTokenService;
    @Autowired
    private OauthRefreshTokenService oauthRefreshTokenService;
    private AuthenticationKeyGenerator authenticationKeyGenerator = new DefaultAuthenticationKeyGenerator();

    @Override
    public OAuth2Authentication readAuthentication(OAuth2AccessToken token) {
        return this.readAuthentication(token.getValue());
    }

    @Override
    public OAuth2Authentication readAuthentication(String token) {
        OAuth2Authentication authentication = null;
        OauthAccessToken oauthAccessToken = new OauthAccessToken();
        oauthAccessToken.setTokenId(extractTokenKey(token));
        oauthAccessToken = oauthAccessTokenService.select(oauthAccessToken);
        try {
            authentication = deserializeAuthentication(oauthAccessToken.getAuthentication());
        } catch (EmptyResultDataAccessException var4) {
            if (LOG.isInfoEnabled()) {
                LOG.info("Failed to find access token for token " + token);
            }
        } catch (IllegalArgumentException var5) {
            LOG.warn("Failed to deserialize authentication for " + token, var5);
            this.removeAccessToken(token);
        }

        return authentication;
    }

    @Override
    public void storeAccessToken(OAuth2AccessToken token, OAuth2Authentication authentication) {
        String refreshToken = null;
        if (token.getRefreshToken() != null) {
            refreshToken = token.getRefreshToken().getValue();
        }
        if (this.readAccessToken(token.getValue()) != null) {
            this.removeAccessToken(token.getValue());
        }
        OauthAccessToken oauthAccessToken = new OauthAccessToken();
        oauthAccessToken.setTokenId(extractTokenKey(token.getValue()));
        oauthAccessToken.setToken(serializeAccessToken(token));
        oauthAccessToken.setId(authenticationKeyGenerator.extractKey(authentication));
        oauthAccessToken.setUserName(authentication.isClientOnly() ? null : authentication.getName());
        oauthAccessToken.setClientId(authentication.getOAuth2Request().getClientId());
        oauthAccessToken.setAuthentication(serializeAuthentication(authentication));
        oauthAccessToken.setRefreshToken(extractTokenKey(refreshToken));
        oauthAccessTokenService.insert(oauthAccessToken);
    }

    public OAuth2AccessToken readAccessToken(String tokenValue) {
        OAuth2AccessToken accessToken = null;
        OauthAccessToken oauthAccessToken = new OauthAccessToken();
        oauthAccessToken.setTokenId(extractTokenKey(tokenValue));
        oauthAccessToken = oauthAccessTokenService.select(oauthAccessToken);
        try {
            accessToken = deserializeAccessToken(oauthAccessToken.getToken());
        } catch (EmptyResultDataAccessException var4) {
            if (LOG.isInfoEnabled()) {
                LOG.info("Failed to find access token for token " + tokenValue);
            }
        } catch (IllegalArgumentException var5) {
            LOG.warn("Failed to deserialize access token for " + tokenValue, var5);
            this.removeAccessToken(tokenValue);
        }

        return accessToken;
    }


    @Override
    public void removeAccessToken(OAuth2AccessToken token) {
        this.removeAccessToken(token.getValue());
    }

    @Override
    public void storeRefreshToken(OAuth2RefreshToken refreshToken, OAuth2Authentication authentication) {
        OauthRefreshToken oauthRefreshToken = new OauthRefreshToken();
        String uuid = UUID.randomUUID().toString().replaceAll("-", "");
        oauthRefreshToken.setId(uuid);
        oauthRefreshToken.setTokenId(extractTokenKey(refreshToken.getValue()));
        oauthRefreshToken.setToken(serializeRefreshToken(refreshToken));
        oauthRefreshToken.setAuthentication(serializeAuthentication(authentication));
        oauthRefreshTokenService.insert(oauthRefreshToken);
    }

    @Override
    public OAuth2RefreshToken readRefreshToken(String token) {
        OAuth2RefreshToken refreshToken = null;
        OauthRefreshToken oauthRefreshToken = oauthRefreshTokenService.findByTokenId(extractTokenKey(token));
        try {
            refreshToken = deserializeRefreshToken(oauthRefreshToken.getToken());
        } catch (EmptyResultDataAccessException var4) {
            if (LOG.isInfoEnabled()) {
                LOG.info("Failed to find refresh token for token " + token);
            }
        } catch (IllegalArgumentException var5) {
            LOG.warn("Failed to deserialize refresh token for token " + token, var5);
            this.removeRefreshToken(token);
        }
        return refreshToken;
    }

    @Override
    public OAuth2Authentication readAuthenticationForRefreshToken(OAuth2RefreshToken token) {
        return readAuthenticationForRefreshToken(token.getValue());
    }

    public OAuth2Authentication readAuthenticationForRefreshToken(String value) {
        OAuth2Authentication authentication = null;
        OauthRefreshToken oauthRefreshToken = oauthRefreshTokenService.findByTokenId(extractTokenKey(value));
        try {
            authentication = deserializeAuthentication(oauthRefreshToken.getAuthentication());
        } catch (EmptyResultDataAccessException var4) {
            if (LOG.isInfoEnabled()) {
                LOG.info("Failed to find access token for token " + value);
            }
        } catch (IllegalArgumentException var5) {
            LOG.warn("Failed to deserialize access token for " + value, var5);
            this.removeRefreshToken(value);
        }
        return authentication;
    }

    @Override
    public void removeRefreshToken(OAuth2RefreshToken token) {
        this.removeRefreshToken(token.getValue());
    }

    public void removeRefreshToken(String token) {
        oauthRefreshTokenService.deleteByTokenId(extractTokenKey(token));
    }

    @Override
    public void removeAccessTokenUsingRefreshToken(OAuth2RefreshToken refreshToken) {
        removeAccessTokenUsingRefreshToken(refreshToken.getValue());
    }

    public void removeAccessTokenUsingRefreshToken(String refreshToken) {
        oauthAccessTokenService.deleteByRefreshToken(extractTokenKey(refreshToken));
    }

    @Override
    public OAuth2AccessToken getAccessToken(OAuth2Authentication authentication) {
        OAuth2AccessToken accessToken = null;
        String key = this.authenticationKeyGenerator.extractKey(authentication);
        OauthAccessToken oauthAccessToken = new OauthAccessToken();
        oauthAccessToken.setId(key);
        oauthAccessToken = oauthAccessTokenService.select(oauthAccessToken);
        try {

            accessToken = deserializeAccessToken(oauthAccessToken.getToken());
        } catch (EmptyResultDataAccessException var5) {
            if (LOG.isDebugEnabled()) {
                LOG.debug("Failed to find access token for authentication " + authentication);
            }
        } catch (IllegalArgumentException var6) {
            LOG.error("Could not extract access token for authentication " + authentication, var6);
        }

        if (accessToken != null && !key.equals(this.authenticationKeyGenerator.extractKey(this.readAuthentication(accessToken.getValue())))) {
            this.removeAccessToken(accessToken.getValue());
            this.storeAccessToken(accessToken, authentication);
        }

        return accessToken;
    }

    @Override
    public Collection<OAuth2AccessToken> findTokensByClientIdAndUserName(String clientId, String userName) {
        List<OAuth2AccessToken> accessTokens = new ArrayList();
        List<OauthAccessToken> oauthAccessTokens = oauthAccessTokenService.findByNameAndClientId(userName,clientId);
        try {
            for(OauthAccessToken oauthAccessToken:oauthAccessTokens){
                OAuth2AccessToken oauth2AccessToken = deserializeAccessToken(oauthAccessToken.getToken());
                accessTokens.add(oauth2AccessToken);
            }
        } catch (EmptyResultDataAccessException var5) {
            if (LOG.isInfoEnabled()) {
                LOG.info("Failed to find access token for clientId " + clientId + " and userName " + userName);
            }
        }
        accessTokens = this.removeNulls((List)accessTokens);
        return accessTokens;
    }
    private List<OAuth2AccessToken> removeNulls(List<OAuth2AccessToken> accessTokens) {
        List<OAuth2AccessToken> tokens = new ArrayList();
        Iterator var3 = accessTokens.iterator();

        while(var3.hasNext()) {
            OAuth2AccessToken token = (OAuth2AccessToken)var3.next();
            if (token != null) {
                tokens.add(token);
            }
        }
        return tokens;
    }
    @Override
    public Collection<OAuth2AccessToken> findTokensByClientId(String clientId) {
        List<OAuth2AccessToken> accessTokens = new ArrayList();
        try {
            List<OauthAccessToken> oauthAccessTokens = oauthAccessTokenService.findByClientId(clientId);
            for(OauthAccessToken oauthAccessToken:oauthAccessTokens){
                accessTokens.add(deserializeAccessToken(oauthAccessToken.getToken()));
            }
        } catch (EmptyResultDataAccessException var4) {
            if (LOG.isInfoEnabled()) {
                LOG.info("Failed to find access token for clientId " + clientId);
            }
        }
        accessTokens = this.removeNulls((List)accessTokens);
        return accessTokens;
    }

    protected byte[] serializeAccessToken(OAuth2AccessToken token) {
        return SerializationUtils.serialize(token);
    }

    protected OAuth2AccessToken deserializeAccessToken(byte[] token) {
        return (OAuth2AccessToken) SerializationUtils.deserialize(token);
    }

    protected OAuth2Authentication deserializeAuthentication(byte[] authentication) {
        return (OAuth2Authentication) SerializationUtils.deserialize(authentication);
    }

    protected byte[] serializeAuthentication(OAuth2Authentication authentication) {
        return SerializationUtils.serialize(authentication);
    }

    protected OAuth2RefreshToken deserializeRefreshToken(byte[] token) {
        return (OAuth2RefreshToken) SerializationUtils.deserialize(token);
    }

    public void removeAccessToken(String tokenValue) {
        oauthAccessTokenService.deleteById(extractTokenKey(tokenValue));
    }

    protected byte[] serializeRefreshToken(OAuth2RefreshToken token) {
        return SerializationUtils.serialize(token);
    }

    protected String extractTokenKey(String value) {
        if (value == null) {
            return null;
        } else {
            MessageDigest digest;
            try {
                digest = MessageDigest.getInstance("MD5");
            } catch (NoSuchAlgorithmException var5) {
                throw new IllegalStateException("MD5 algorithm not available.  Fatal (should be in the JDK).");
            }

            try {
                byte[] bytes = digest.digest(value.getBytes("UTF-8"));
                return String.format("%032x", new BigInteger(1, bytes));
            } catch (UnsupportedEncodingException var4) {
                throw new IllegalStateException("UTF-8 encoding not available.  Fatal (should be in the JDK).");
            }
        }
    }
}
